/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"

#include <iostream>

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = -1;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
  RC rc;

  // Open page file
  if ((rc = pf.open(indexname, mode))) 
    return rc;

  // Load headers from page 0
  char buf[PageFile::PAGE_SIZE];
  if ((rc = pf.read(0, buf))) {
    // If rc is RC_INVALID_PID, then this pagefile was just created (new index)
    if (rc == RC_INVALID_PID) {
      rootPid = -1;
      treeHeight = 0;
      if ((rc = flushHeaders()))
        return rc;
      return 0;
    }
    return rc;
  }

  rootPid = *((PageId *) buf);
  treeHeight = *((int *) (buf + sizeof(PageId)));

  return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
  return pf.close();
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
  RC rc;

  // If the index is empty, create a single leaf node
  if (rootPid == -1) {
    // Build leaf node
    BTLeafNode btln;
    btln.insert(key, rid);
    btln.setNextNodePtr(-1);

    // Commit it to pagefile
    if ((rc = btln.write(1, pf)))
      return rc;

    // Fill in fields
    rootPid = 1;
    treeHeight = 1;
 
    if ((rc = flushHeaders()))
      return rc;

    return 0;
  }

  // Insert into a non-empty index
  bool newChild = false;
  int childKey;
  PageId childPid;
  if ((rc = _insert(key, rid, rootPid, treeHeight - 1, newChild, childKey, childPid)))
    return rc;

  // If a new child node is created, we have to make a new root node
  if (newChild) {
    BTNonLeafNode btnln;
    PageId newRootPid = pf.endPid();

    // Create new root node
    if ((rc = btnln.initializeRoot(rootPid, childKey, childPid)))
      return rc;

    // Commit to pagefile and update rootPid and treeHeight
    if ((rc = btnln.write(newRootPid, pf)))
      return rc;
    rootPid = newRootPid;
    treeHeight++;

    if ((rc = flushHeaders()))
      return rc;
  }
  
  return 0;
}

/*
 * Recusive implementation of insert
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @param pid[IN] the current page in the tree
 * @param depth[IN] the current depth in the tree
 * @param newChild[OUT] whether a new child node was created
 * @param childKey[OUT] the child key of the new node
 * @param childPid[OUT] the pid of the new node
 * @return error code. 0 if no error
 */
RC BTreeIndex::_insert(int key, const RecordId& rid, const PageId& pid, int depth, bool& newChild, int& childKey, PageId& childPid) 
{
  RC rc;

  // If depth is 0, we are at the leaf level
  if (depth == 0) {
    BTLeafNode btln;
    
    // Load the leaf node page
    if ((rc = btln.read(pid, pf)))
      return rc;

    // Attempt to insert the key into the node.
    // If the return code is non-zero, we need to split and insert
    if (btln.insert(key, rid)) {
      // Get a new page for the sibling node.
      // No need to do anything with the sibling buffer since it will be
      // overwritten anyways.
      BTLeafNode sibling;
      PageId siblingPid = pf.endPid();

      // Insert and split leaf node
      int siblingKey;
      if ((rc = btln.insertAndSplit(key, rid, sibling, siblingKey)))
        return rc;
      btln.setNextNodePtr(siblingPid);

      // Write both nodes to the pagefile
      if ((rc = btln.write(pid, pf)))
        return rc;
      if ((rc = sibling.write(siblingPid, pf)))
        return rc;

      // Tell the parent node to add a new (key, pid) pair
      newChild = true;
      childKey = siblingKey;
      childPid = siblingPid;
      
      return 0;
    }

    // Write out node to pagefile
    if ((rc = btln.write(pid, pf)))
      return rc;

    // No new child
    newChild = false;
    
    return 0;
  }

  // We are in a non-leaf node.

  // Find where next node at a lower level in the tree
  BTNonLeafNode btnln;
  PageId nextPid;
  if ((rc = btnln.read(pid, pf)))
    return rc;
  if ((rc = btnln.locateChildPtr(key, nextPid)))
    return rc;

  bool atEnd;
  if (nextPid == btnln.pLastPid[0])
    atEnd = true;
  else
    atEnd = false;

  // Recurse
  bool _newChild = false;
  int _childKey;
  PageId _childPid;
  if ((rc = _insert(key, rid, nextPid, depth - 1, _newChild, _childKey, _childPid)))
    return rc;

  // Add new (key, pid) pair to this node if necessary
  if (_newChild) {
    // Try to insert the pair into this node
    if (btnln.insert(_childKey, _childPid, atEnd)) {
      // Split required
      // Create sibling node
      BTNonLeafNode sibling;
      PageId siblingPid = pf.endPid();

      // Insert and split
      newChild = true;
      btnln.insertAndSplit(_childKey, _childPid, sibling, childKey);
      childPid = siblingPid;

      // Write this node and sibling out to the pagefile
      if ((rc = btnln.write(pid, pf)))
        return rc;
      if ((rc = sibling.write(siblingPid, pf)))
        return rc;

      return 0;
    }

/*    if (depth == 1) {
      print();
      int x;
      cin >> x;
    }*/

    newChild = false;

    // Write this node to pagefile
    rc = btnln.write(pid, pf);
    return rc;
  }

  // No new child. Do nothing.
  newChild = false;

  return 0;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
    return _locate(searchKey, cursor, rootPid, treeHeight - 1);;
}

/**
 * Recursive implementation of locate
 * @param searchKey[IN] the key to search for
 * @param cursor[OUT] cursor containing the location of the key pair
 * @param pid[IN] the current pid
 * @param depth[IN] the current depth in the tree
 */
RC BTreeIndex::_locate(int searchKey, IndexCursor& cursor, const PageId& pid, int depth) 
{
  RC rc;

  // If depth is 0, we are at the leaf node.
  if (depth == 0) {
    // Load the leaf node
    BTLeafNode btln;
    if ((rc = btln.read(pid, pf)))
      return rc;

    // Find the entry
    int eid, key;
    RecordId rid;
    if ((rc = btln.locate(searchKey, eid)))
      return rc;
    if ((rc = btln.readEntry(eid, key, rid)))
      return rc;

    // Set the fields in cursor
    cursor.pid = pid;
    cursor.eid = eid;
    return 0;
  }

  // We are in a non-leaf node, search for the next node down the tree
  // Load the non-leaf node
  BTNonLeafNode btnln;
  btnln.read(pid, pf);

  // Find the page containing the next node
  PageId nextPid;
  if ((rc = btnln.locateChildPtr(searchKey, nextPid)))
    return rc;

  // Recurse
  return _locate(searchKey, cursor, nextPid, depth - 1);
}


/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
  // Make sure we have a valid pid in the cursor
  PageId endPid = pf.endPid();
  if (cursor.pid <= 0 || cursor.pid >= endPid)
    return RC_INVALID_CURSOR;

  // Read the leaf node
  BTLeafNode btln;
  btln.read(cursor.pid, pf);

  // Read the entry
  btln.readEntry(cursor.eid, key, rid);

  // Move the cursor forward
  // If we can move forward within the node
  if (cursor.eid < btln.getKeyCount() - 1) {
    cursor.eid++;
  }
  // Otherwise we need to move across nodes
  else {
    PageId nextPid = btln.getNextNodePtr();
    // If there are no more entries
    if (nextPid == -1) {
      // Invalidate the cursor
      cursor.pid = -1;
      cursor.eid = -1;
    }
    // Otherwise set cursor to first element of the next node
    else {
      cursor.pid = nextPid;
      cursor.eid = 0;
    }
  }

  return 0;
}

// Write header page to disk
RC BTreeIndex::flushHeaders() {
  char buf[PageFile::PAGE_SIZE];
  *((PageId *) buf) = rootPid;
  *((int *) (buf + sizeof(PageId))) = treeHeight;

  return pf.write(0, buf);
}
